from mcp.server.fastmcp import FastMCP
import math
import secrets
import random
from typing import Union, Optional

# 创建整合后的MCP服务实例
mcp = FastMCP("Big Number Operations Service")

# 四则运算
@mcp.tool()
def big_number_add(a: str, b: str) -> str:
    """大数加法运算
    Args:
        a: 第一个大数字符串
        b: 第二个大数字符串
    Returns:
        两数之和的字符串表示
    """
    try:
        num_a = int(a)
        num_b = int(b)
        return str(num_a + num_b)
    except ValueError:
        return "错误：输入不是有效的整数"

@mcp.tool()
def big_number_subtract(a: str, b: str) -> str:
    """大数减法运算
    Args:
        a: 被减数大数字符串
        b: 减数大数字符串
    Returns:
        两数之差的字符串表示
    """
    try:
        num_a = int(a)
        num_b = int(b)
        return str(num_a - num_b)
    except ValueError:
        return "错误：输入不是有效的整数"

@mcp.tool()
def big_number_multiply(a: str, b: str) -> str:
    """大数乘法运算
    Args:
        a: 第一个大数字符串
        b: 第二个大数字符串
    Returns:
        两数之积的字符串表示
    """
    try:
        num_a = int(a)
        num_b = int(b)
        return str(num_a * num_b)
    except ValueError:
        return "错误：输入不是有效的整数"

@mcp.tool()
def big_number_divide(a: str, b: str) -> str:
    """大数除法运算（整数除法）
    Args:
        a: 被除数大数字符串
        b: 除数大数字符串(不能为0)
    Returns:
        两数之商的字符串表示（整数除法）
    """
    try:
        num_a = int(a)
        num_b = int(b)
        if num_b == 0:
            return "错误:除数不能为0"
        return str(num_a // num_b)
    except ValueError:
        return "错误：输入不是有效的整数"

# 高级运算
@mcp.tool()
def compute_gcd(a: int, b: int) -> int:
    """计算两个大数的最大公约数
    Args:
        a: 第一个整数
        b: 第二个整数
    Returns:
        两个数的最大公约数
    """
    if not (isinstance(a, int) and isinstance(b, int)):
        raise ValueError("输入必须为整数")
    a = abs(a)
    b = abs(b)
    if a == 0 and b == 0:
        raise ValueError("两个数不能同时为0")
    return math.gcd(a, b)

@mcp.tool()
def compute_lcm(a: int, b: int) -> int:
    """计算两个大整数的最小公倍数
    Args:
        a: 第一个整数
        b: 第二个整数
    Returns:
        两个数的最小公倍数
    """
    if a == 0 or b == 0:
        raise ValueError("输入不能包含0")
    gcd = math.gcd(a, b)
    return abs(a * b) // gcd

@mcp.tool()
def power(base: int, exponent: int) -> int:
    """大数乘方运算
    Args:
        base: 底数
        exponent: 指数（非负整数）
    Returns:
        底数的指数次幂
    """
    if exponent < 0:
        raise ValueError("指数必须为非负整数")
    return base ** exponent

@mcp.tool()
def square_root(n: int) -> Union[int, float]:
    """大数开平方运算（返回整数部分或精确浮点数）
    Args:
        n: 被开方数（非负整数）
    Returns:
        平方根（整数部分或精确浮点数）
    """
    if n < 0:
        raise ValueError("被开方数必须为非负整数")
    root = math.isqrt(n)
    if root * root == n:
        return root
    return math.sqrt(n)

@mcp.tool()
def logarithm(n: int, base: int = 10) -> float:
    """大数对数运算（以base为底）
    Args:
        n: 真数（正整数）
        base: 底数（正整数，默认10）
    Returns:
        对数值（浮点数）
    """
    if n <= 0 or base <= 0:
        raise ValueError("真数和底数必须为正整数")
    return math.log(n, base)

@mcp.tool()
def compare(a: int, b: int) -> str:
    """大数比较运算
    Args:
        a: 第一个数
        b: 第二个数
    Returns:
        比较结果（'a > b'/'a == b'/'a < b'）
    """
    if a > b:
        return 'a > b'
    elif a == b:
        return 'a == b'
    else:
        return 'a < b'

@mcp.tool()
def modular_inverse(a: int, m: int) -> Optional[int]:
    """计算大数模逆运算
    
    Args:
        a: 需要求逆元的整数
        m: 模数
        
    Returns:
        返回a在模m下的逆元，如果不存在则返回None
    """
    try:
        # 处理边界情况
        if m <= 1:
            return None
            
        # 处理负数情况
        a = a % m
        
        # 检查是否互质
        if math.gcd(a, m) != 1:
            return None
            
        # 使用扩展欧几里得算法计算模逆
        def extended_gcd(a, b):
            if b == 0:
                return (a, 1, 0)
            else:
                # 使用迭代方式避免递归深度限制
                old_r, r = a, b
                old_s, s = 1, 0
                old_t, t = 0, 1
                
                while r != 0:
                    quotient = old_r // r
                    old_r, r = r, old_r - quotient * r
                    old_s, s = s, old_s - quotient * s
                    old_t, t = t, old_t - quotient * t
                
                return (old_r, old_s, old_t)
                
        g, x, _ = extended_gcd(a, m)
        if g != 1:
            return None
        else:
            result = x % m
            # 验证结果是否正确
            if (a * result) % m == 1:
                return result
            return None
            
    except Exception:
        return None

# 椭圆曲线运算
@mcp.tool()
def point_addition(p: tuple, q: tuple, a: int, p_mod: int) -> tuple:
    """椭圆曲线点加法运算
    
    Args:
        p: 第一个点 (x1, y1)
        q: 第二个点 (x2, y2)
        a: 椭圆曲线参数a
        p_mod: 模数
        
    Returns:
        点加法结果 (x3, y3)
    """
    try:
        x1, y1 = p
        x2, y2 = q
        
        # 处理无穷远点情况
        if x1 is None and y1 is None:
            return q
        if x2 is None and y2 is None:
            return p
        
        # 处理点互为逆元情况
        if x1 == x2 and (y1 + y2) % p_mod == 0:
            return (None, None)
        
        # 计算斜率
        if p == q:
            # 点加倍情况
            numerator = (3 * x1 * x1 + a) % p_mod
            denominator = (2 * y1) % p_mod
        else:
            # 不同点相加
            numerator = (y2 - y1) % p_mod
            denominator = (x2 - x1) % p_mod
        
        # 计算分母的模逆
        inv_denominator = pow(denominator, -1, p_mod)
        s = (numerator * inv_denominator) % p_mod
        
        # 计算新点坐标
        x3 = (s * s - x1 - x2) % p_mod
        y3 = (s * (x1 - x3) - y1) % p_mod
        
        return (x3, y3)
    except Exception as e:
        raise ValueError(f"点加法运算失败: {str(e)}")

@mcp.tool()
def point_multiplication(k: int, p: tuple, a: int, p_mod: int) -> tuple:
    """椭圆曲线点乘法运算
    
    Args:
        k: 乘数
        p: 基点 (x, y)
        a: 椭圆曲线参数a
        p_mod: 模数
        
    Returns:
        点乘法结果 (x, y)
    """
    try:
        if k == 0:
            return (None, None)
            
        result = (None, None)
        current = p
        
        while k > 0:
            if k % 2 == 1:
                result = point_addition(result, current, a, p_mod)
            current = point_addition(current, current, a, p_mod)
            k = k // 2
            
        return result
    except Exception as e:
        raise ValueError(f"点乘法运算失败: {str(e)}")

# 位运算
@mcp.tool()
def bitwise_and(a: int, b: int) -> int:
    """位与运算
    Args:
        a: 第一个整数
        b: 第二个整数
    Returns:
        两数的位与结果
    """
    return a & b

@mcp.tool()
def bitwise_or(a: int, b: int) -> int:
    """位或运算
    Args:
        a: 第一个整数
        b: 第二个整数
    Returns:
        两数的位或结果
    """
    return a | b

@mcp.tool()
def bitwise_xor(a: int, b: int) -> int:
    """位异或运算
    Args:
        a: 第一个整数
        b: 第二个整数
    Returns:
        两数的位异或结果
    """
    return a ^ b

@mcp.tool()
def bitwise_not(a: int) -> int:
    """位非运算
    Args:
        a: 输入整数
    Returns:
        位非结果
    """
    return ~a

@mcp.tool()
def left_shift(a: int, shift: int) -> int:
    """左移运算
    Args:
        a: 输入整数
        shift: 位移位数
    Returns:
        左移结果
    """
    return a << shift

@mcp.tool()
def right_shift(a: int, shift: int) -> int:
    """右移运算
    Args:
        a: 输入整数
        shift: 位移位数
    Returns:
        右移结果
    """
    return a >> shift

# 随机数生成
@mcp.tool()
def generate_large_number(bits: int) -> int:
    """随机生成指定二进制位数的大整数
    Args:
        bits: 生成的大整数的二进制位数（至少1位）
    Returns:
        指定二进制位数的大整数
    """
    if bits < 1:
        raise ValueError("位数不能小于1")
    # 生成安全随机数，最高位为1以保证位数准确
    return secrets.randbits(bits) | (1 << (bits - 1))

@mcp.tool()
def generate_large_prime(bits: int = 1024) -> int:
    """生成指定位数的大素数
    Args:
        bits: 素数位数（默认1024位）
    Returns:
        大素数
    """
    if bits < 2:
        raise ValueError("素数位数至少为2位")
        
    while True:
        candidate = secrets.randbits(bits)
        candidate |= (1 << (bits - 1)) | 1  # 确保最高位和最低位为1
        if is_prime(candidate):
            return candidate

# 模运算
@mcp.tool()
def modular_exponentiation(base: int, exponent: int, modulus: int) -> int:
    """大数模幂运算
    Args:
        base: 底数
        exponent: 指数
        modulus: 模数
    Returns:
        (base^exponent) mod modulus
    """
    if modulus == 1:
        return 0
    result = 1
    base = base % modulus
    while exponent > 0:
        if exponent % 2 == 1:
            result = (result * base) % modulus
        exponent = exponent >> 1
        base = (base * base) % modulus
    return result

# 素性检测
@mcp.tool()
def check_prime(number: Union[int, str], rounds: int = 10) -> str:
    """检测一个大整数是否为素数
    Args:
        number: 待检测的整数（可以是整数或数字字符串）
        rounds: 检测轮数，轮数越多准确率越高，默认10轮
    Returns:
        str: 检测结果，包含是否为素数的判断和置信度
    """
    try:
        # 验证输入
        if isinstance(number, str):
            if not number.isdigit():
                return f"错误：输入'{number}'不是有效的整数"
            number = int(number)
            
        if number < 0:
            return "错误：不支持负数检测"
            
        # 执行素性检测
        result = is_prime(number, rounds)
        confidence = (1 - (1/4)**rounds) * 100
        
        if result:
            return f"{number} 是素数（置信度：{confidence:.6f}%）"
        else:
            return f"{number} 不是素数"
    except Exception as e:
        return f"检测失败：{str(e)}"

def is_prime(n: Union[int, str], k: int = 10) -> bool:
    """使用Miller-Rabin算法检测大整数是否为素数
    Args:
        n: 待检测的整数（可以是整数或数字字符串）
        k: 检测轮数，轮数越多准确率越高
    Returns:
        bool: 是否为素数
    """
    # 处理字符串输入
    if isinstance(n, str):
        n = int(n)
        
    # 处理小数字的情况
    if n <= 1:
        return False
    elif n <= 3:
        return True
    elif n % 2 == 0:
        return False
    
    # 将n-1分解为d*2^s
    d = n - 1
    s = 0
    while d % 2 == 0:
        d //= 2
        s += 1
    
    # Miller-Rabin测试
    for _ in range(k):
        a = secrets.randbelow(n - 2) + 2  # 使用更安全的随机数生成
        x = pow(a, d, n)  # 计算a^d mod n
        
        if x == 1 or x == n - 1:
            continue
        
        for __ in range(s - 1):
            x = pow(x, 2, n)
            if x == n - 1:
                break
        else:
            # 如果经过s-1次平方后仍不等于n-1，则n是合数
            return False
    
    # 通过所有测试，n很可能是素数
    return True

if __name__ == "__main__":
    print("大数运算综合服务启动...")
    mcp.run(transport="stdio")